#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import json
import hashlib
import time
from typing import Dict, List, Any, Optional

import numpy as np
from pymilvus import MilvusClient
from pymilvus.orm.types import DataType

from vector_store.vector_store_interface import VectorStoreInterface

# 获取logger
logger = logging.getLogger(__name__)

class MilvusVectorStore(VectorStoreInterface):
    """
    使用Milvus作为向量数据库，存储知识图谱节点的向量表示
    """

    def __init__(self, 
                 uri: str = None, 
                 token: str = None,
                 db_path: str = None,
                 collection_name: str = None,
                 dimension: int = 1024):
        """初始化Milvus向量存储

        Args:
            uri: Milvus服务URI，如果为None则使用Milvus Lite
            token: Milvus服务认证令牌
            db_path: Milvus Lite数据库文件路径，仅在uri为None时使用
            collection_name: 集合名称
            dimension: 向量维度
        """
        self.uri = uri
        self.token = token
        self.db_path = db_path
        self.collection_name = collection_name
        self.dimension = dimension
        self.vid_cache = {}  # 缓存已经生成的VID，避免重复计算
        self.collection_loaded = False  # 新增：跟踪集合加载状态
        
        # 初始化Milvus客户端
        if uri:
            self.client = MilvusClient(uri=uri, token=token)
            logger.info(f"已连接到Milvus服务: {uri}")
        else:
            self.client = MilvusClient(db_path)
            logger.info(f"已初始化Milvus Lite: {db_path}")
        
        # 初始化集合
        self.initialize_collection()

    def initialize_collection(self) -> None:
        """初始化Milvus集合，创建必要的结构"""
        try:
            # 检查集合是否已存在
            if self.client.has_collection(self.collection_name):
                logger.info(f"集合已存在: {self.collection_name}")
                
                # 尝试获取集合信息，检查维度是否匹配
                try:
                    collection_info = self.client.describe_collection(self.collection_name)
                    logger.info(f"集合信息: {collection_info}")
                    
                    # 检查维度是否匹配
                    if 'schema' in collection_info and 'fields' in collection_info['schema']:
                        for field in collection_info['schema']['fields']:
                            if field.get('name') == 'vector':
                                field_dim = field.get('params', {}).get('dim', 0)
                                if field_dim != self.dimension:
                                    logger.warning(f"集合维度不匹配: 期望 {self.dimension}, 实际 {field_dim}。将删除并重新创建集合")
                                    self.client.drop_collection(self.collection_name)
                                    logger.info(f"旧集合已删除，准备重新创建")
                                    break  # 集合已删除，将继续执行下面的创建流程
                except Exception as e:
                    logger.warning(f"获取集合信息失败: {str(e)}")
                
                # 如果集合仍然存在，尝试加载到内存
                if self.client.has_collection(self.collection_name):
                    try:
                        self.client.load_collection(self.collection_name)
                        logger.info(f"集合已加载到内存: {self.collection_name}")
                        self.collection_loaded = True  # 设置加载状态标记
                        return  # 如果成功加载集合，直接返回
                    except Exception as load_error:
                        logger.error(f"加载集合失败: {str(load_error)}")
                        # 如果加载失败，而且是严重错误，则尝试重新创建集合
                        try:
                            self.client.drop_collection(self.collection_name)
                            logger.info(f"旧集合已删除，准备重新创建")
                            # 集合已删除，将继续执行下面的创建流程
                        except Exception as drop_error:
                            logger.error(f"删除集合失败: {str(drop_error)}")
                            raise drop_error
            
            # 到这里说明集合不存在或已被删除，需要创建新集合
            # 创建集合
            logger.info(f"创建集合: {self.collection_name}，向量维度: {self.dimension}")
            
            # 创建schema对象
            schema = self.client.create_schema()
            
            # 添加ID字段，使用VARCHAR类型
            schema.add_field(
                field_name="id", 
                datatype=DataType.VARCHAR, 
                is_primary=True, 
                max_length=36
            )
            
            # 添加其他字段
            schema.add_field(field_name="original_id", datatype=DataType.VARCHAR, max_length=256)
            schema.add_field(field_name="node_type", datatype=DataType.VARCHAR, max_length=32)
            schema.add_field(field_name="node_text", datatype=DataType.VARCHAR, max_length=10000)
            schema.add_field(field_name="metadata", datatype=DataType.VARCHAR, max_length=10000)
            schema.add_field(field_name="created_at", datatype=DataType.INT64)
            schema.add_field(field_name="updated_at", datatype=DataType.INT64)
            
            # 添加向量字段
            schema.add_field(
                field_name="vector", 
                datatype=DataType.FLOAT_VECTOR, 
                dim=self.dimension
            )
            
            # 创建集合
            self.client.create_collection(
                collection_name=self.collection_name,
                schema=schema,
                consistency_level="Strong"
            )
            
            # 为向量字段创建索引 - 集合创建后必须创建索引
            try:
                # 创建索引参数对象
                index_params = self.client.prepare_index_params()
                index_params.add_index(
                    field_name="vector",
                    index_type="HNSW",
                    metric_type="COSINE",
                    params={"M": 8, "efConstruction": 200}
                )
                
                # 应用索引
                self.client.create_index(
                    collection_name=self.collection_name,
                    index_params=index_params
                )
                logger.info(f"向量索引创建成功")
            except Exception as index_error:
                logger.error(f"创建索引失败: {str(index_error)}")
                # 索引创建失败是严重错误，因为没有索引将无法进行向量搜索
                # 但我们依然尝试加载集合，可能某些基本查询仍然可用
            
            # 加载集合到内存
            try:
                self.client.load_collection(self.collection_name)
                logger.info(f"集合已加载到内存: {self.collection_name}")
                self.collection_loaded = True  # 设置加载状态标记
            except Exception as load_error:
                logger.error(f"加载集合到内存失败: {str(load_error)}")
                # 加载失败是严重错误，但我们不抛出异常，让后续操作继续尝试
            
            logger.info(f"Milvus集合初始化成功: {self.collection_name}")
        except Exception as e:
            logger.error(f"Milvus集合初始化失败: {str(e)}")
            raise

    def _generate_vid(self, original_id: str) -> str:
        """
        生成32位VID字符串，与图数据库中使用的ID保持一致

        Args:
            original_id: 原始ID字符串

        Returns:
            32位固定长度的字符串
        """
        # 使用缓存避免重复计算相同ID
        if original_id in self.vid_cache:
            return self.vid_cache[original_id]

        # 使用MD5生成32位字符串，确保符合图数据库VID长度要求
        vid = hashlib.md5(original_id.encode('utf-8')).hexdigest()

        # 记录到缓存
        self.vid_cache[original_id] = vid
        return vid

    def ensure_collection_loaded(self) -> bool:
        """确保集合已加载到内存中，如果需要则进行加载"""
        try:
            # 如果已经标记为加载状态，直接返回
            if self.collection_loaded:
                logger.debug(f"集合已加载状态，无需重新加载: {self.collection_name}")
                return True
                
            # 检查集合是否存在
            if not self.client.has_collection(self.collection_name):
                logger.warning(f"集合不存在，尝试初始化: {self.collection_name}")
                self.initialize_collection()
                return True
            
            # 检查集合是否已加载
            collection_info = self.client.describe_collection(self.collection_name)
            logger.debug(f"集合信息: {collection_info}")
            
            # 尝试加载集合，不再检查索引是否存在
            logger.info(f"集合未加载，正在加载: {self.collection_name}")
            try:
                # 直接尝试加载集合
                self.client.load_collection(self.collection_name)
                logger.info(f"集合已成功加载: {self.collection_name}")
                
                # 设置加载状态标记
                self.collection_loaded = True
                return True
            except Exception as load_error:
                logger.error(f"加载集合失败: {str(load_error)}")
                return False
        except Exception as e:
            logger.error(f"检查集合加载状态时出错: {str(e)}")
            return False

    def store_node_vector(self, vid: str, node_type: str, node_text: str, vector: List[float], metadata: Dict[str, Any] = None) -> bool:
        """
        存储节点向量

        Args:
            vid: 节点ID
            node_type: 节点类型 ('class', 'function' 或 'annotation')
            node_text: 节点文本内容
            vector: 节点的嵌入向量
            metadata: 节点元数据

        Returns:
            操作是否成功
        """
        try:
            if node_type not in ['class', 'function', 'annotation']:
                logger.warning(f"不支持的节点类型: {node_type}，仅支持class、function和annotation")
                return False

            # 确保集合已加载
            if not self.ensure_collection_loaded():
                logger.error(f"存储节点向量失败: 集合未加载")
                return False

            # 验证向量维度并处理
            vector_len = len(vector)
            if vector_len != self.dimension:
                logger.warning(f"向量维度不匹配: 期望 {self.dimension}, 实际 {vector_len}，尝试调整")
                
                # 如果向量维度大于期望，截断
                if vector_len > self.dimension:
                    logger.info(f"截断向量维度 {vector_len} -> {self.dimension}")
                    vector = vector[:self.dimension]
                # 如果向量维度小于期望，填充0
                else:
                    logger.info(f"填充向量维度 {vector_len} -> {self.dimension}")
                    vector = vector + [0.0] * (self.dimension - vector_len)

            # 生成与图数据库一致的VID
            nebula_vid = self._generate_vid(vid)
            logger.debug(f"原始ID: {vid} -> 向量存储VID: {nebula_vid}")

            # 准备元数据
            meta_str = json.dumps(metadata) if metadata else "{}"
            
            # 准备时间戳
            current_time = int(time.time())
            
            # 准备插入数据 - 确保ID是字符串类型
            data = {
                "id": str(nebula_vid),  # 确保是字符串类型
                "vector": vector,  # 向量字段
                "original_id": str(vid),
                "node_type": node_type,
                "node_text": node_text,
                "metadata": meta_str,
                "created_at": current_time,
                "updated_at": current_time
            }
            
            # 检查是否已存在，如果存在则先删除
            exists = self.get_node_vector(vid)
            if exists:
                self.delete_node_vector(vid)
            
            # 插入数据
            result = self.client.insert(
                collection_name=self.collection_name,
                data=[data]
            )
            
            logger.info(f"节点向量已存储/更新: {vid} ({node_type})")
            return True
        except Exception as e:
            logger.error(f"存储节点向量失败: {str(e)}")
            return False

    def get_node_vector(self, vid: str) -> Optional[Dict[str, Any]]:
        """
        获取节点向量

        Args:
            vid: 节点ID

        Returns:
            节点数据字典或None
        """
        try:
            # 确保集合已加载
            if not self.ensure_collection_loaded():
                logger.error(f"获取节点向量失败: 集合未加载")
                return None
            
            # 生成与图数据库一致的VID
            nebula_vid = self._generate_vid(vid)
            
            # 查询数据 - 确保ID是字符串类型
            result = self.client.query(
                collection_name=self.collection_name,
                ids=[str(nebula_vid)],
                output_fields=["original_id", "node_type", "node_text", "vector", "metadata", "created_at", "updated_at"]
            )
            
            # 检查结果
            if not result or len(result) == 0:
                return None
            
            # 处理结果
            entity = result[0]
            
            return {
                'vid': nebula_vid,
                'original_id': entity.get('original_id'),
                'node_type': entity.get('node_type'),
                'node_text': entity.get('node_text'),
                'vector': entity.get('vector'),
                'metadata': json.loads(entity.get('metadata', '{}')),
                'created_at': entity.get('created_at'),
                'updated_at': entity.get('updated_at')
            }
        except Exception as e:
            logger.error(f"获取节点向量失败: {str(e)}")
            return None

    def get_node_by_original_id(self, original_id: str) -> Optional[Dict[str, Any]]:
        """
        通过原始ID获取节点向量

        Args:
            original_id: 原始节点ID

        Returns:
            节点数据字典或None
        """
        try:
            # 确保集合已加载
            if not self.ensure_collection_loaded():
                logger.error(f"通过原始ID获取节点向量失败: 集合未加载")
                return None
            
            # 查询数据
            result = self.client.query(
                collection_name=self.collection_name,
                filter=f"original_id == '{original_id}'",
                output_fields=["id", "original_id", "node_type", "node_text", "vector", "metadata", "created_at", "updated_at"]
            )
            
            # 检查结果
            if not result or len(result) == 0:
                return None
            
            # 处理结果
            entity = result[0]
            
            return {
                'vid': entity.get('id'),
                'original_id': entity.get('original_id'),
                'node_type': entity.get('node_type'),
                'node_text': entity.get('node_text'),
                'vector': entity.get('vector'),
                'metadata': json.loads(entity.get('metadata', '{}')),
                'created_at': entity.get('created_at'),
                'updated_at': entity.get('updated_at')
            }
        except Exception as e:
            logger.error(f"通过原始ID获取节点向量失败: {str(e)}")
            return None

    def search_similar_vectors(self, query_vector: List[float], node_type: Optional[str] = None, top_k: int = 10) -> List[Dict[str, Any]]:
        """
        搜索与查询向量最相似的节点

        Args:
            query_vector: 查询向量
            node_type: 节点类型过滤 ('class', 'function' 或 'annotation')
            top_k: 返回结果数量

        Returns:
            相似节点列表，按相似度降序排序
        """
        try:
            # 确保集合已加载
            if not self.ensure_collection_loaded():
                logger.error(f"搜索向量失败: 集合未加载")
                return []
                
            # 验证向量维度并处理
            vector_len = len(query_vector)
            if vector_len != self.dimension:
                logger.warning(f"查询向量维度不匹配: 期望 {self.dimension}, 实际 {vector_len}，尝试调整")
                
                # 如果向量维度大于期望，截断
                if vector_len > self.dimension:
                    logger.info(f"截断查询向量维度 {vector_len} -> {self.dimension}")
                    query_vector = query_vector[:self.dimension]
                # 如果向量维度小于期望，填充0
                else:
                    logger.info(f"填充查询向量维度 {vector_len} -> {self.dimension}")
                    query_vector = query_vector + [0.0] * (self.dimension - vector_len)
                
            # 准备过滤条件
            filter_expr = None
            if node_type:
                if node_type not in ['class', 'function', 'annotation']:
                    logger.warning(f"不支持的节点类型: {node_type}，仅支持class、function和annotation")
                    return []
                filter_expr = f"node_type == '{node_type}'"
            
            # 执行向量搜索
            search_results = self.client.search(
                collection_name=self.collection_name,
                data=[query_vector],
                filter=filter_expr,
                limit=top_k,
                output_fields=["original_id", "node_type", "node_text", "metadata", "created_at", "updated_at"]
            )
            
            # 处理搜索结果
            if not search_results or len(search_results) == 0:
                return []
            
            # 解析结果 - Milvus返回的是一个列表，每个元素对应一个查询向量的结果
            results_list = []
            for hit in search_results[0]:
                entity = hit.get('entity', {})
                results_list.append({
                    'vid': str(hit.get('id')),  # 确保ID是字符串类型
                    'original_id': entity.get('original_id'),
                    'node_type': entity.get('node_type'),
                    'node_text': entity.get('node_text'),
                    'metadata': json.loads(entity.get('metadata', '{}')),
                    'similarity': 1 - hit.get('distance', 0),  # 转换余弦距离为相似度
                    'created_at': entity.get('created_at'),
                    'updated_at': entity.get('updated_at')
                })
            
            return results_list
        except Exception as e:
            logger.error(f"搜索向量失败: {str(e)}")
            return []

    def delete_node_vector(self, vid: str) -> bool:
        """
        删除节点向量

        Args:
            vid: 节点ID

        Returns:
            操作是否成功
        """
        try:
            # 确保集合已加载
            if not self.ensure_collection_loaded():
                logger.error(f"删除节点向量失败: 集合未加载")
                return False
                
            # 生成与图数据库一致的VID
            nebula_vid = self._generate_vid(vid)
            
            # 删除数据 - 确保ID是字符串类型
            result = self.client.delete(
                collection_name=self.collection_name,
                ids=[str(nebula_vid)]
            )
            
            # 检查结果
            if not result or str(nebula_vid) not in result:
                logger.warning(f"删除节点向量失败，可能不存在: {vid}")
                return False
            
            logger.info(f"节点向量已删除: {vid}")
            return True
        except Exception as e:
            logger.error(f"删除节点向量失败: {str(e)}")
            return False

    def get_stats(self) -> Dict[str, Any]:
        """
        获取向量存储统计信息

        Returns:
            统计信息字典
        """
        try:
            # 获取集合统计信息
            # Milvus没有直接提供集合统计的API，我们需要自己查询
            
            # 获取实体总数
            total_count = 0
            try:
                # 查询所有实体，但只获取ID字段以减少带宽
                result = self.client.query(
                    collection_name=self.collection_name,
                    output_fields=["id"]
                )
                total_count = len(result) if result else 0
            except Exception as e:
                logger.warning(f"获取实体总数失败: {str(e)}")
            
            # 获取各类型节点数量
            type_counts = {}
            for node_type in ['class', 'function', 'annotation']:
                try:
                    result = self.client.query(
                        collection_name=self.collection_name,
                        filter=f"node_type == '{node_type}'",
                        output_fields=["id"]
                    )
                    type_counts[node_type] = len(result) if result else 0
                except Exception as e:
                    logger.warning(f"获取{node_type}类型节点数量失败: {str(e)}")
                    type_counts[node_type] = 0
            
            return {
                'total': total_count,
                'by_type': type_counts,
                'collection_name': self.collection_name,
                'dimension': self.dimension,
                'store_type': 'milvus'
            }
        except Exception as e:
            logger.error(f"获取向量存储统计信息失败: {str(e)}")
            return {
                'total': 0,
                'by_type': {'class': 0, 'function': 0, 'annotation': 0},
                'collection_name': self.collection_name,
                'dimension': self.dimension,
                'store_type': 'milvus',
                'error': str(e)
            }

    def close(self) -> None:
        """关闭资源"""
        # Milvus客户端不需要显式关闭
        logger.info("Milvus向量存储已关闭") 